#!/usr/bin/python
# coding:utf-8
import tflearn
from tflearn.layers.core import input_data, fully_connected
from tflearn.layers.conv import conv_2d, max_pool_2d
from tflearn.layers.estimator import regression
import tflearn.datasets.mnist as mnist
import tensorflow as tf

# 读取数据
trainX, trainY, testX, testY = mnist.load_data("../le_datasets/mnist/input_data",one_hot=True)  # todo 填写mnist地址

trainX = trainX.reshape([-1, 28, 28, 1])
testX = testX.reshape([-1, 28, 28, 1])  # todo shape 中的-1代表什么
                                        # 定义的时候, 这个维度不确定, 就用-1表示

# 构建网络
net = input_data(shape=[None, 28, 28, 1], name="input")  # input封装了placeholder来接入输入数据

net = conv_2d(net, 32, 5, activation="relu")  # 深度为32, 尺寸为5*5, 激活函数为relu的卷积层
net = max_pool_2d(net, 2)  # 过滤器尺寸为2 的最大池化层

net = conv_2d(net, 64, 5, activation="relu")
net = max_pool_2d(net, 2)

net = fully_connected(net, 500, activation="relu")
net = fully_connected(net, 10, activation="softmax")

# 封装好的函数学习任务. 这里指定了优化器为sgd, 学习率为0.01, 损失函数为交叉熵
net = regression(net, optimizer="sgd", learning_rate=0.01, loss="categorical_crossentropy")

# tflearn 将训练过程封装到一个类中
model = tflearn.DNN(net, tensorboard_verbose=0)  # TODO 猜测这里的参数设置为1时的样子
model.fit(trainX, trainY, n_epoch=20, validation_set=([testX,testY]), show_metric=True)

writer = tf.summary.FileWriter("../log/tflearn", tf.get_default_graph())
writer.close()